#! /usr/bin/env python3

# 1.导包
import rospy
from superglue.srv import SuperScore, SuperScoreRequest, SuperScoreResponse
from cv_bridge import CvBridge
import cv_bridge
print(' - cv_bridge.__file__ = ',cv_bridge.__file__)
from sensor_msgs.msg import Image
import sensor_msgs


import cv2
import numpy as np
import torch
import yaml


import sys
sys.path.append("/home/daybeha/Documents/github/DeepLabV3_ws/src/superglue")
from Detectors import create_detector
from Matchers import create_matcher

from models.utils import AverageTimer
from utils.tools import *


class Matching(torch.nn.Module):
    """ Image Matching Frontend (SuperPoint + SuperGlue) """
    def __init__(self, config={}):
        super().__init__()
        # create detector
        self.detector = create_detector(config["detector"])
        # create matcher
        self.matcher = create_matcher(config["matcher"])


        # self.superpoint = SuperPoint(config.get('superpoint', {}))
        # self.superglue = SuperGlue(config.get('superglue', {}))

    def forward(self, data):
        """ Run SuperPoint (optionally) and SuperGlue
        SuperPoint is skipped if ['keypoints0', 'keypoints1'] exist in input
        Args:
          data: dictionary with minimal keys: ['image0', 'image1']
        """
        pred = {'ref': None, 'cur': None}

        # TODO 这块的显存存占用有待优化
        # Extract SuperPoint (keypoints, scores, descriptors) if not provided
        if 'keypoints0' not in data:
            pred['ref'] = self.detector(data['image0'])
        if 'keypoints1' not in data:
            pred['cur'] = self.detector(data['image1'])

        matches = self.matcher(pred)

        return matches



class Server():
    def __init__(self):
        config = "/home/daybeha/Documents/github/DeepLabV3_ws/src/superglue/params/superpoint_supergluematch.yaml"
        with open(config, 'r') as f:
            config = yaml.safe_load(f)
        # print(f"cuda avaliable: {torch.cuda.is_available()}")
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print("device: ", device)
        config["device"] = device
        self.matching = Matching(config).eval().to(device)
        self.bridge = CvBridge()

        self.imgs = []
        # 3.创建服务对象
        self.server = rospy.Service("/super_score", SuperScore, self.imgCallback)
        self.wait = False
        self.wait_t = rospy.Rate(0.1)

        self.cnt = 0

    def compute_score(self, img0, img1):
        matches = self.matching({'image0': img0, 'image1': img1})
        score = np.mean(matches["match_score"].cpu().detach().numpy())
        num = matches["match_score"].shape[0]
        # print(f"score: {score}   num: {num}")
        return score, num

    def imgCallback(self, req):
        rospy.loginfo(f"received request {self.cnt}")
        self.cnt +=1
        # img0 = bridge.imgmsg_to_cv2(req.image0, "bgr8")
        # img1 = bridge.imgmsg_to_cv2(req.image1, "bgr8")
        img0 = self.bridge.imgmsg_to_cv2(req.image0, "mono8")
        img1 = self.bridge.imgmsg_to_cv2(req.image1, "mono8")

        # self.imgs.append(img0, img1)

        while(self.wait):
            rospy.logerr("waitiing")
            self.wait_t.sleep()

        self.wait = True
        score, num = self.compute_score(img0, img1)
        self.wait = False

        resp_score = SuperScoreResponse()
        resp_score.score = score
        resp_score.match_num = num
        rospy.loginfo(f"send response: {score}\t match num:{num}\n")
        return resp_score

def con_show(img0, img1):
    # 纵向连接 image = np.vstack((img0, img1))
    # 横向连接 image = np.concatenate([img0, img1], axis=1)
    image = np.concatenate((img0, img1))

    cv2.imshow("image in python", image)
    # cv2.imwrite("/home/tt/test/111.png", img)
    cv2.waitKey(0)


def show_img(img=None):
    cv2.imshow("image in python", img)
    # cv2.imwrite("/home/tt/test/111.png", img)
    cv2.waitKey(1)




if __name__ == "__main__":
    # 这一句至关重要！！！ 能节省至少一半显存！！！
    torch.set_grad_enabled(False)

    # 2.初始化 ROS 节点
    rospy.init_node("superglue_server")

    # 4.回调函数处理请求并产生响应
    server = Server()

    # 5.spin 函数
    rospy.spin()





